#pragma once

#include "BasicOp.h"
#include "ParamsLoader.h"
#include "Structure.h"

typedef DL::Tensor1D<float> T1;
typedef DL::Tensor2D<float> T2;
typedef DL::Tensor3D<float> T3;
typedef DL::Tensor4D<float> T4;

namespace TEST {
	bool IsEqual(float v1, float v2) {
		float diff = v1 - v2;
		diff = diff > 0 ? diff : -diff;
		//if (diff < 0.00001f)//0.000015
		if (v1 == v2)
		{
			return true;
		}
		else
		{
			printf("error find %f,between %f -%f\r\n", diff, v1, v2);
			return false;
		}
	}
	bool CompareTensor1D(const DL::Tensor1D<float>& t1, const DL::Tensor1D<float>& t2);
	bool CompareTensor2D(const DL::Tensor2D<float>& t1, const DL::Tensor2D<float>& t2);
	bool CompareTensor3D(const DL::Tensor3D<float>& t1, const DL::Tensor3D<float>& t2);

	void DoTest(bool success, std::string name) {
		printf("Test %s %s\r\n", name.c_str(), success ? "success" : "failed");
	}

	bool TestLinear(std::string testBenchFile)
	{
		ParamsLoader loader;
		loader.LoadFromFile(testBenchFile);
		const T1& input = loader.LoadTensor1D("input");
		const T2& weight = loader.LoadTensor2D("weight");
		const T1& bias = loader.LoadTensor1D("bias");
		const T1& output = loader.LoadTensor1D("output");


		T1& result = DL::Linear(input, weight, bias);
		bool success = CompareTensor1D(output, result);
		delete& result;
		return success;
	}

	bool TestConv2D(std::string testBenchFile)
	{
		ParamsLoader loader;
		loader.LoadFromFile(testBenchFile);
		const T3& input = loader.LoadTensor3D("input");
		const T4& weight = loader.LoadTensor4D("weight");
		const T1& bias = loader.LoadTensor1D("bias");
		const T3& output = loader.LoadTensor3D("output");
		int stride = loader.LoadIntValue("stride");
		int padding = loader.LoadIntValue("padding");

		T3& result = DL::Conv2D(input, stride, padding, weight, bias);
		bool success = CompareTensor3D(output, result);
		delete& result;
		return success;
	}
	bool TestRelu(std::string testBenchFile)
	{
		ParamsLoader loader;
		loader.LoadFromFile(testBenchFile);
		const T1& input = loader.LoadTensor1D("input");
		const T1& output = loader.LoadTensor1D("output");

		T1& result = DL::Relu(input);
		bool success = CompareTensor1D(output, result);
		delete& result;
		return success;
	}
	bool TestMaxPooling(std::string testBenchFile)
	{
		ParamsLoader loader;
		loader.LoadFromFile(testBenchFile);
		const T3& input = loader.LoadTensor3D("input");
		const T3& output = loader.LoadTensor3D("output");
		int kernalSize = loader.LoadIntValue("kernalSize");
		int stride = loader.LoadIntValue("stride");
		int padding = loader.LoadIntValue("padding");

		T3& result = DL::MaxPooling2D(input, kernalSize, stride, padding);
		bool success = CompareTensor3D(output, result);
		delete& result;
		return success;
	}
	bool TestFlatten(std::string testBenchFile)
	{
		ParamsLoader loader;
		loader.LoadFromFile(testBenchFile);
		const T3& input = loader.LoadTensor3D("input");
		const T1& output = loader.LoadTensor1D("output");

		T1& result = DL::Flatten(input);
		bool success = CompareTensor1D(output, result);
		delete& result;
		return success;
	}
	bool TestNormalize(std::string imageMark, std::string testBenchFile) {
		Image image(imageMark);
		DL::Tensor3D<float> imageTensor(image);
		float mean[] = { 0.5f,0.5f,0.5f };
		float std[] = { 0.5f,0.5f,0.5f };
		DL::Tensor3D<float>& normalized = DL::Normalize(imageTensor, mean, std);
		ParamsLoader loader;
		loader.LoadFromFile(testBenchFile);
		const DL::Tensor3D<float>& torch = loader.LoadTensor3D("normalized");
		bool success = TEST::CompareTensor3D(normalized, torch);
		delete& normalized;
		return success;
	}

	bool CompareTensor1D(const DL::Tensor1D<float>& t1, const DL::Tensor1D<float>& t2)
	{
		bool equal = true;
		for (int i = 0; i < t1.width; i++)
		{
			float v1 = t1.Get(i);
			float v2 = t2.Get(i);
			if (!IsEqual(v1, v2)) equal = false;
		}
		return equal;
	}
	bool CompareTensor2D(const DL::Tensor2D<float>& t1, const  DL::Tensor2D<float>& t2)
	{
		bool equal = true;
		for (int y = 0; y < t1.height; y++)
		{
			for (int x = 0; x < t1.width; x++)
			{
				float v1 = t1.Get(x, y);
				float v2 = t2.Get(x, y);
				if (!IsEqual(v1, v2)) equal = false;
			}
		}
		return equal;
	}
	bool CompareTensor3D(const DL::Tensor3D<float>& t1, const DL::Tensor3D<float>& t2)
	{
		bool equal = true;
		for (int c = 0; c < t1.chanel; c++)
		{
			for (int y = 0; y < t1.height; y++)
			{
				for (int x = 0; x < t1.width; x++)
				{
					float v1 = t1.Get(x, y, c);
					float v2 = t2.Get(x, y, c);
					if (!IsEqual(v1, v2)) equal = false;
				}
			}
		}
		return equal;
	}
}